import dtshare as dt2
from extern import *
import abc
import multiprocessing.pool
import gc
class fetch_stock_data(metaclass=abc.ABCMeta):
    def __init__(self,stock_list,date):
        self.stock_list=stock_list
        self.date=date
    def get_stocks_by_range(self,code):
        df =dt2.stock_zh_a_tick(code,self.date)
        df.insert(0,'trade_date',self.date)
        #df.insert(0,'code',code)
        return({'code':code,'df':df})
    def fetch_data(self, stock_list):
        """获取股票信息"""
        pool = multiprocessing.pool.ThreadPool(len(stock_list))
        try:
            res = pool.map(self.get_stocks_by_range, stock_list)
        finally:
            pool.close()
        return [d for d in res if d is not None]
    def getdata(self):
        return(self.fetch_data(self.stock_list))
class writetodbs(metaclass=abc.ABCMeta):
    def __init__(self,data,date):
        self.data=data
        self.date=date
        con,db=mongocon()
        db=con['deals']
        self.con=con
        self.db=db
    def writetodb(self,df):
        try:
            code=df['code']
            df=df['df']
            if df.size>0:
                tmp = json.loads(df.T.to_json()).values()
                self.db[code].delete_many({'trade_date':self.date})
                self.db[code].insert_many(tmp)
            re=0
        except:
            re=code
        return(re)
    def write_data(self, stock_list):
        """获取股票信息"""
        pool = multiprocessing.pool.ThreadPool(len(stock_list))
        try:
            res = pool.map(self.writetodb, stock_list)
        finally:
            pool.close()
            self.con.close()
        return [d for d in res if d !=0]
    def writedata(self):
        return(self.write_data(self.data))
def qiege(stocks,lens=100):
    if len(stocks) % lens==0:
        nums=int(len(stocks)/lens)
    else:
        nums = int(len(stocks)/lens)+1
    new=[]
    for i in range(nums):
        tmp=stocks[i*lens:(i+1)*lens]
        new.append(tmp)
    return(new)
def update(date):
    stocks =getstocks()
    newstocks=[]
    for i in stocks:
        if i[:3]=='688' or i[:3]=='300':continue #过滤科创板股票
        newstocks.append(i[-2:].lower()+i[:6])
    newstocks=sorted(newstocks)
    print('newstocks:',len(newstocks))
    tmpfile='update_s.txt'   #记录未写入成功的股票
    if os.path.exists(tmpfile):
        with open(tmpfile,'r')as f:
            last = f.readline()
            index = newstocks.index(last)
            print('last is ',last,'index',index)
            newstocks=newstocks[index+1:]
    for i in qiege(newstocks):
        dd = fetch_stock_data(i,date)
        start = time.clock()
        data = dd.getdata()
        end=time.clock()
        print('get data',len(data),round(end-start,2))
        ww=writetodbs(data,date)
        start = time.clock()
        ww = ww.writedata()
        end=time.clock()
        print('write data',len(ww),round(end-start,2))
        if len(ww)>0:
            for j in ww:
                print(j,'not write over')
                writetofile('update.log',j+' '+date+' did not write over\n')
        with open(tmpfile,'w')as f:
            f.write(i[-1])
        gc.collect()
    remove(tmpfile)
if __name__=='__main__':
    start ='20200301'
    end = latestday()
    tradedays=trade_days(start,end)
    datefile='date.txt'
    if os.path.exists(datefile):
        with open(datefile,'r')as f:
            last = f.readline()
            print('last date is',last)
    else:last=''
    for i in range(len(tradedays)):
        if last!='':
            if i<=tradedays.index(last):continue
        i=tradedays[i]
        print(i)
        update(i)
        with open(datefile,'w')as f:
            f.write(i)
    #os.system('net stop mongodb')
    #os.system('shutdown -s')
    
